import os
from dotenv import load_dotenv
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_deepseek import ChatDeepSeek

load_dotenv(verbose=True)

api_key = os.getenv("DEEPSEEK_API_KEY")
api_base = os.getenv("DEEPSEEK_API_BASE")


llm = ChatDeepSeek(
    temperature=0.5,
    model="deepseek-chat",
    api_key=api_key,
    api_base=api_base
)

parser = JsonOutputParser()

# 自定义 JSON 格式说明，明确要求中文回答
format_instructions = """
你必须以 JSON 格式回答，并且必须使用中文。
你的回答应该是一个包含以下字段的 JSON 对象:
{
  "笑话": "这里是笑话内容（必须用中文）",
  "结语": "这里是笑话的结语或笑点（必须用中文）"
}
注意：回答必须是中文，不要使用英文。
"""

prompt_template = """
{format_instructions}

用中文回答以下问题: {query}
"""

prompt = PromptTemplate(
    template=prompt_template,
    input_variables=["query"],
    partial_variables={"format_instructions": format_instructions}
)
print(prompt.format(query="请给我讲一个笑话"))

chain = prompt | llm | parser

result = chain.invoke({"query": "请给我讲一个笑话"})

print(result)
